Diagonal traverse II

Time: O(MxN); Space: O(M); medium

Given a list of lists of integers, nums, return all elements of nums in diagonal order as shown in the below images.

Example 1:

Input: nums = [[1,2,3],[4,5,6],[7,8,9]]

Output: [1,4,2,7,5,3,8,6,9]

Example 2:

Input: nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]

Output: [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

Example 3:

Input: nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]

Output: [1,4,2,5,3,8,6,9,7,10,11]

Example 4:

Input: nums = [[1,2,3,4,5,6]]

Output: [1,2,3,4,5,6]

Constraints:

  • 1 <= len(nums) <= 10^5

  • 1 <= len(nums[i]) <= 10^5

  • 1 <= nums[i][j] <= 10^9

  • There at most 10^5 elements in nums.

Hints:

  1. Notice that numbers with equal sums of row and column indexes belong to the same diagonal.

  2. Store them in tuples (sum, row, val), sort them, and then regroup the answer.

[1]:
import collections

class Solution1(object):
    """
    Time: O(M*N)
    Space: O(M)
    """
    def findDiagonalOrder(self, nums):
        """
        :type nums: List[List[int]]
        :rtype: List[int]
        """
        result, dq, col = [], collections.deque(), 0

        for i in range(len(nums) + max(list(map(len, nums)))-1):
            new_dq = collections.deque()
            if i < len(nums):
                dq.appendleft((i, 0))
            for r, c in dq:
                result.append(nums[r][c])
                if c+1 < len(nums[r]):
                    new_dq.append((r, c+1))
            dq = new_dq

        return result
[2]:
s = Solution1()

nums = [[1,2,3],[4,5,6],[7,8,9]]
assert s.findDiagonalOrder(nums) == [1,4,2,7,5,3,8,6,9]

nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
assert s.findDiagonalOrder(nums) == [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
assert s.findDiagonalOrder(nums) == [1,4,2,5,3,8,6,9,7,10,11]

nums = [[1,2,3,4,5,6]]
assert s.findDiagonalOrder(nums) == [1,2,3,4,5,6]
[3]:
class Solution2(object):
    def findDiagonalOrder(self, nums):
        """
        :type nums: List[List[int]]
        :rtype: List[int]
        """
        result = []
        for r, row in enumerate(nums):
            for c, num in enumerate(row):
                if len(result) <= r+c:
                    result.append([])
                result[r+c].append(num)

        return [num for row in result for num in reversed(row)]
[4]:
s = Solution2()

nums = [[1,2,3],[4,5,6],[7,8,9]]
assert s.findDiagonalOrder(nums) == [1,4,2,7,5,3,8,6,9]

nums = [[1,2,3,4,5],[6,7],[8],[9,10,11],[12,13,14,15,16]]
assert s.findDiagonalOrder(nums) == [1,6,2,8,7,3,9,4,12,10,5,13,11,14,15,16]

nums = [[1,2,3],[4],[5,6,7],[8],[9,10,11]]
assert s.findDiagonalOrder(nums) == [1,4,2,5,3,8,6,9,7,10,11]

nums = [[1,2,3,4,5,6]]
assert s.findDiagonalOrder(nums) == [1,2,3,4,5,6]